/**
 * @file test_eigen_basics.cpp
 *
 * Test the basic add-ons for Eigen library
 * 
 * @author Dahua Lin
 */

#include <gtest/gtest.h>

#include <slipp/base/arrays.h>

#include <functional>
#include <iostream>

#include "eigen_test.h"

using namespace sli;
using namespace sli::test;

template<typename T>
struct doubler : public std::unary_function<T, T>
{
	T operator() (const T& x) const
	{
		return x * T(2);
	}
};

darow_t forloop_transform(const darow_t& x, double (*f)(double) )
{
	index_t n = x.cols();
	darow_t y(n);

	for (index_t i = 0; i < n; ++i)
	{
		y(i) = f(x(i));
	}

	return y;
}

darow_t forloop_transform(const darow_t& x1, const darow_t& x2, double (*f)(double, double) )
{
	index_t n = x1.cols();
	darow_t y(n);

	for (index_t i = 0; i < n; ++i)
	{
		y(i) = f(x1(i), x2(i));
	}

	return y;
}


TEST( EigenCommon, BasicMakers )
{
	double src[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

	darr_t r1 = make_arow(6, src);
	EXPECT_EQ( r1.rows(), 1 );
	EXPECT_EQ( r1.cols(), 6 );
	EXPECT_EQ( r1(0), 1.0 );
	EXPECT_EQ( r1(1), 2.0 );
	EXPECT_EQ( r1(2), 3.0 );
	EXPECT_EQ( r1(3), 4.0 );
	EXPECT_EQ( r1(4), 5.0 );
	EXPECT_EQ( r1(5), 6.0 );

	darr_t r2 = make_acol(6, src);
	EXPECT_EQ( r2.rows(), 6 );
	EXPECT_EQ( r2.cols(), 1 );
	EXPECT_EQ( r2(0), 1.0 );
	EXPECT_EQ( r2(1), 2.0 );
	EXPECT_EQ( r2(2), 3.0 );
	EXPECT_EQ( r2(3), 4.0 );
	EXPECT_EQ( r2(4), 5.0 );
	EXPECT_EQ( r2(5), 6.0 );

	darr_t r3 = make_arr(2, 3, src);
	EXPECT_EQ( r3.rows(), 2 );
	EXPECT_EQ( r3.cols(), 3 );
	EXPECT_EQ( r3(0, 0), 1.0 );
	EXPECT_EQ( r3(1, 0), 2.0 );
	EXPECT_EQ( r3(0, 1), 3.0 );
	EXPECT_EQ( r3(1, 1), 4.0 );
	EXPECT_EQ( r3(0, 2), 5.0 );
	EXPECT_EQ( r3(1, 2), 6.0 );

	dmat_t r4 = make_vrow(6, src);
	EXPECT_EQ( r4.rows(), 1 );
	EXPECT_EQ( r4.cols(), 6 );
	EXPECT_EQ( r4(0), 1.0 );
	EXPECT_EQ( r4(1), 2.0 );
	EXPECT_EQ( r4(2), 3.0 );
	EXPECT_EQ( r4(3), 4.0 );
	EXPECT_EQ( r4(4), 5.0 );
	EXPECT_EQ( r4(5), 6.0 );

	dmat_t r5 = make_vcol(6, src);
	EXPECT_EQ( r5.rows(), 6 );
	EXPECT_EQ( r5.cols(), 1 );
	EXPECT_EQ( r5(0), 1.0 );
	EXPECT_EQ( r5(1), 2.0 );
	EXPECT_EQ( r5(2), 3.0 );
	EXPECT_EQ( r5(3), 4.0 );
	EXPECT_EQ( r5(4), 5.0 );
	EXPECT_EQ( r5(5), 6.0 );

	darr_t r6 = make_mat(2, 3, src);
	EXPECT_EQ( r6.rows(), 2 );
	EXPECT_EQ( r6.cols(), 3 );
	EXPECT_EQ( r6(0, 0), 1.0 );
	EXPECT_EQ( r6(1, 0), 2.0 );
	EXPECT_EQ( r6(0, 1), 3.0 );
	EXPECT_EQ( r6(1, 1), 4.0 );
	EXPECT_EQ( r6(0, 2), 5.0 );
	EXPECT_EQ( r6(1, 2), 6.0 );
}


TEST( EigenCommon, Mapping )
{
	darow_t a = darow_t::LinSpaced(5, 1, 5);
	darow_t b = darow_t::LinSpaced(5, 3, 15);

	darow_t r = map(doubler<double>(), a);

	ASSERT_EQ( r.rows(), 1 );
	ASSERT_EQ( r.cols(), 5 );
	EXPECT_TRUE( test_equal(r, darow_t::LinSpaced(5, 2, 10)) );

	darow_t r2 = map(std::plus<double>(), a, b);

	ASSERT_EQ( r2.rows(), 1 );
	ASSERT_EQ( r2.cols(), 5 );
	EXPECT_TRUE( test_equal(r2, darow_t::LinSpaced(5, 4, 20)) );
}


TEST( EigenCommon, Filtering )
{
	bool bsrc[6] = {true, false, true, true, false, false};

	int32_t r1_g[3] = {0, 2, 3};
	iacol_t r1 = find( make_arow(6, bsrc) );
	ASSERT_EQ( r1.rows(), 3 );
	ASSERT_EQ( r1.cols(), 1 );
	EXPECT_TRUE( test_equal(r1, make_acol(3, r1_g)) );

	double xsrc[6] = {3, 4, 5, 1, 2, 6};
	double ysrc[6] = {4, 5, 1, 3, 2, 4};

	int32_t r2_g[3] = {1, 2, 5};
	iacol_t r2 = find( std::bind2nd(std::greater<double>(), 3), make_arow(6, xsrc) );
	ASSERT_EQ( r2.rows(), 3 );
	ASSERT_EQ( r2.cols(), 1 );
	EXPECT_TRUE( test_equal(r2, make_acol(3, r2_g)) );

	int32_t r3_g[4] = {0, 1, 3, 4};
	iacol_t r3 = find( std::less_equal<double>(), make_arow(6, xsrc), make_arow(6, ysrc));
	ASSERT_EQ( r3.rows(), 4 );
	ASSERT_EQ( r3.cols(), 1 );
	EXPECT_TRUE( test_equal(r3, make_acol(4, r3_g)) );
}


TEST( EigenCommon, SelectCoeff )
{
	// array 1d

	iarr_t idx1(2, 3); idx1 << 3, 7, 2, 4, 1, 3;
	Eigen::Array4i idx2; idx2 << 3, 7, 2, 4;

	darow_t a(8);  a << 0, 2, 4, 6, 8, 10, 12, 14;

	darr_t ra1g(2, 3); ra1g << 6, 14, 4, 8, 2, 6;
	dacol_t ra2g(4, 1); ra2g << 6, 14, 4, 8;

	darr_t ra1 = select_coeffs(a, idx1);
	Eigen::Array4d ra2 = select_coeffs(a, idx2);

	ASSERT_EQ( ra1.rows(), 2 );
	ASSERT_EQ( ra1.cols(), 3 );
	EXPECT_TRUE( test_equal(ra1, ra1g) );

	ASSERT_EQ( ra2.rows(), 4 );
	ASSERT_EQ( ra2.cols(), 1 );
	EXPECT_TRUE( test_equal(ra2, ra2g) );

	// matrix 1d

	dvrow_t b(8);  b << 0, 2, 4, 6, 8, 10, 12, 14;

	dmat_t rb1 = select_coeffs(b, idx1);
	Eigen::Vector4d rb2 = select_coeffs(b, idx2);

	ASSERT_EQ( rb1.rows(), 2 );
	ASSERT_EQ( rb1.cols(), 3 );
	EXPECT_TRUE( test_equal(rb1, ra1g) );

	ASSERT_EQ( rb2.rows(), 4 );
	ASSERT_EQ( rb2.cols(), 1 );
	EXPECT_TRUE( test_equal(ra2, ra2g) );

	// 2d

	iacol_t irow(4); irow << 0, 1, 2, 1;
	iacol_t icol(4); icol << 2, 0, 1, 0;

	darr_t A(3, 3);
	A << 1, 2, 3,
		 4, 5, 6,
		 7, 8, 9;

	dmat_t B(3, 3);
	B << 1, 2, 3,
		 4, 5, 6,
		 7, 8, 9;

	dacol_t rA = select_coeffs(A, irow, icol);
	dacol_t rA0(4); rA0 << 3, 4, 8, 4;

	ASSERT_EQ( rA.rows(), 4 );
	ASSERT_EQ( rA.cols(), 1 );
	EXPECT_TRUE( test_equal(rA, rA0) );

	dvcol_t rB = select_coeffs(B, irow, icol);
	dvcol_t rB0(4); rB0 << 3, 4, 8, 4;

	ASSERT_EQ( rB.rows(), 4 );
	ASSERT_EQ( rB.cols(), 1 );
	EXPECT_TRUE( test_equal(rB, rB0) );
}


TEST( EigenCommon, SelectRowsCols )
{
	darr_t A(4, 4);
	A <<  1,  2,  3,  4,
		  5,  6,  7,  8,
		  9, 10, 11, 12,
		 13, 14, 15, 16;

	dmat_t M(4, 4);
	M <<  1,  2,  3,  4,
		  5,  6,  7,  8,
		  9, 10, 11, 12,
		 13, 14, 15, 16;

	// select rows

	iarow_t irows(2); irows << 2, 1;

	darr_t sr0(2, 4);
	sr0 << 9, 10, 11, 12,
		   5, 6, 7, 8;

	darr_t A_sr = select_rows(A, irows);
	dmat_t M_sr = select_rows(M, irows);

	ASSERT_EQ( A_sr.rows(), 2 );
	ASSERT_EQ( A_sr.cols(), 4 );
	ASSERT_EQ( M_sr.rows(), 2 );
	ASSERT_EQ( M_sr.cols(), 4 );

	EXPECT_TRUE( test_equal(A_sr, sr0) );
	EXPECT_TRUE( test_equal(M_sr, sr0) );

	// select cols

	iarow_t icols(3); icols << 3, 0, 1;

	darr_t sc0(4, 3);
	sc0 << 4, 1, 2,
		   8, 5, 6,
		   12, 9, 10,
		   16, 13, 14;

	darr_t A_sc = select_cols(A, icols);
	dmat_t M_sc = select_cols(M, icols);

	ASSERT_EQ( A_sc.rows(), 4 );
	ASSERT_EQ( A_sc.cols(), 3 );
	ASSERT_EQ( M_sc.rows(), 4 );
	ASSERT_EQ( M_sc.cols(), 3 );

	EXPECT_TRUE( test_equal(A_sc, sc0) );
	EXPECT_TRUE( test_equal(M_sc, sc0) );

	// select rows & cols

	darr_t src0(2, 3);
	src0 << 12, 9, 10,
		    8, 5, 6;

	darr_t A_src = select_rows_cols(A, irows, icols);
	dmat_t M_src = select_rows_cols(M, irows, icols);

	ASSERT_EQ( A_src.rows(), 2 );
	ASSERT_EQ( A_src.cols(), 3 );
	ASSERT_EQ( M_src.rows(), 2 );
	ASSERT_EQ( M_src.cols(), 3 );

	EXPECT_TRUE( test_equal(A_src, src0) );
	EXPECT_TRUE( test_equal(M_src, src0) );

}


TEST( EigenShare, ShareArrays )
{
	share_darow_t r = make_share_arow<double>(3);

	ASSERT_EQ( r.rows(), 1 );
	ASSERT_EQ( r.cols(), 3 );

	share_dacol_t c = make_share_acol<double>(4);

	ASSERT_EQ( c.rows(), 4 );
	ASSERT_EQ( c.cols(), 1 );

	share_darr_t a = make_share_arr<double>(3, 4);

	ASSERT_EQ( a.rows(), 3 );
	ASSERT_EQ( a.cols(), 4 );

	share_darr_t a2 = a;

	ASSERT_EQ( a2.rows(), 3 );
	ASSERT_EQ( a2.cols(), 4 );
	ASSERT_EQ( a.data(), a2.data() );

	share_darr_t b = share_copy(a.get());

	ASSERT_EQ( a.data(), a2.data() );
	ASSERT_NE( a.data(), b.data() );
}


TEST( EigenShare, ShareMatrices )
{
	share_dvrow_t r = make_share_vrow<double>(3);

	ASSERT_EQ( r.rows(), 1 );
	ASSERT_EQ( r.cols(), 3 );

	share_dvcol_t c = make_share_vcol<double>(4);

	ASSERT_EQ( c.rows(), 4 );
	ASSERT_EQ( c.cols(), 1 );

	share_dmat_t a = make_share_mat<double>(3, 4);

	ASSERT_EQ( a.rows(), 3 );
	ASSERT_EQ( a.cols(), 4 );

	share_dmat_t a2 = a;

	ASSERT_EQ( a2.rows(), 3 );
	ASSERT_EQ( a2.cols(), 4 );
	ASSERT_EQ( a.data(), a2.data() );

	share_dmat_t b = share_copy(a.get());

	ASSERT_EQ( a.data(), a2.data() );
	ASSERT_NE( a.data(), b.data() );
}



TEST( EigenEFuns, PowerFuns )
{
	darow_t a = darow_t::LinSpaced(9, -1, 1);
	darow_t c2 = darow_t::Constant(9, 2.0);

	darow_t r_abs = vabs(a);
	ASSERT_EQ( r_abs.rows(), 1 );
	ASSERT_EQ( r_abs.cols(), 9 );
	EXPECT_TRUE( test_equal( r_abs, forloop_transform(a, math::abs) ) );

	EXPECT_TRUE( test_equal( vsqr(a), forloop_transform(a, math::sqr) ) );

	EXPECT_TRUE( test_approx( vsqrt(a), forloop_transform(a, math::sqrt), 1.0e-15 ) );
	EXPECT_TRUE( test_approx( vcube(a), forloop_transform(a, math::cube), 1.0e-15 ) );
	EXPECT_TRUE( test_approx( vrcp(a),  forloop_transform(a, math::rcp),  1.0e-15 ) );

	EXPECT_TRUE( test_approx( vpow(a, 2.0),  forloop_transform(a, math::sqr),  1.0e-14 ) );
	EXPECT_TRUE( test_approx( vpow(a, 0.5),  forloop_transform(a, math::sqrt),  1.0e-14 ) );
	EXPECT_TRUE( test_approx( vpow(a, 3.0),  forloop_transform(a, math::cube),  1.0e-14 ) );
	EXPECT_TRUE( test_approx( vpow(a, -1.0),  forloop_transform(a, math::rcp),  1.0e-14 ) );

	EXPECT_TRUE( test_approx( vpow(a, c2),  forloop_transform(a, math::sqr),  1.0e-14) );
}


TEST( EigenEFuns, ExpLogFuns )
{
	darow_t a = darow_t::LinSpaced(20, 0.1, 2.0);

	EXPECT_TRUE( test_equal( vexp(a), forloop_transform(a, math::exp) ) );
	EXPECT_TRUE( test_equal( vlog(a), forloop_transform(a, math::log) ) );
}


TEST( EigenEFuns, RoundFuns )
{
	darow_t a = darow_t::LinSpaced(9, -1, 1);

	EXPECT_TRUE( test_equal( vfloor(a), forloop_transform(a, math::floor) ) );
	EXPECT_TRUE( test_equal( vceil(a), forloop_transform(a, math::ceil) ) );
	EXPECT_TRUE( test_equal( vround(a), forloop_transform(a, math::round) ) );
}


TEST( EigenEFuns, TrigFuns )
{
	darow_t a = darow_t::LinSpaced(9, -1, 1);
	darow_t b = darow_t::LinSpaced(9, 0, 1);

	EXPECT_TRUE( test_approx(vsin(a),  forloop_transform(a, math::sin), 1.0e-14) );
	EXPECT_TRUE( test_approx(vcos(a),  forloop_transform(a, math::cos),  1.0e-14) );
	EXPECT_TRUE( test_approx(vtan(a),  forloop_transform(a, math::tan),  1.0e-14) );

	EXPECT_TRUE( test_approx(vasin(a),  forloop_transform(a, math::asin), 1.0e-14) );
	EXPECT_TRUE( test_approx(vacos(a),  forloop_transform(a, math::acos), 1.0e-14) );
	EXPECT_TRUE( test_approx(vatan(a),  forloop_transform(a, math::atan), 1.0e-14) );

	EXPECT_TRUE( test_approx(vatan2(a, b),  forloop_transform(a, b, math::atan2), 1.0e-14) );
}


TEST( EigenEFuns, HyperFuns )
{
	darow_t a = darow_t::LinSpaced(9, -1, 1);

	EXPECT_TRUE( test_approx(vsinh(a),  forloop_transform(a, math::sinh), 1.0e-14) );
	EXPECT_TRUE( test_approx(vcosh(a),  forloop_transform(a, math::cosh),  1.0e-14) );
	EXPECT_TRUE( test_approx(vtanh(a),  forloop_transform(a, math::tanh),  1.0e-14) );
}


TEST( EigenAlgorithms, InplaceSort )
{
	darow_t a1(5);
	a1 << 3, 4, 1, 9, 2;

	darow_t a2 = a1;

	inplace_sort_asc(a1);
	inplace_sort_des(a2);

	ASSERT_EQ( a1.size(), 5 );
	ASSERT_EQ( a2.size(), 5 );

	darow_t r1(5);
	r1 << 1, 2, 3, 4, 9;

	darow_t r2(5);
	r2 << 9, 4, 3, 2, 1;

	EXPECT_TRUE( test_equal(a1, r1) );
	EXPECT_TRUE( test_equal(a2, r2) );
}


TEST( EigenAlgorithms, Sort )
{
	darow_t a(5);
	a << 3, 4, 1, 9, 2;

	darow_t b1, b2;

	sort_asc(a, b1);
	sort_des(a, b2);

	ASSERT_EQ( b1.size(), 5 );
	ASSERT_EQ( b2.size(), 5 );

	darow_t r1(5);
	r1 << 1, 2, 3, 4, 9;

	darow_t r2(5);
	r2 << 9, 4, 3, 2, 1;

	EXPECT_TRUE( test_equal(b1, r1) );
	EXPECT_TRUE( test_equal(b2, r2) );
}


TEST( EigenAlgorithms, SortWithIndex )
{
	darow_t a(5);
	a << 3, 4, 1, 9, 2;

	darow_t b1, b2;
	iarow_t i1, i2;

	sort_with_index_asc(a, b1, i1);
	sort_with_index_des(a, b2, i2);

	ASSERT_EQ( b1.size(), 5 );
	ASSERT_EQ( i1.size(), 5 );
	ASSERT_EQ( b2.size(), 5 );
	ASSERT_EQ( i2.size(), 5 );

	darow_t r1(5); r1 << 1, 2, 3, 4, 9;
	iarow_t ri1(5); ri1 << 2, 4, 0, 1, 3;

	darow_t r2(5); r2 << 9, 4, 3, 2, 1;
	iarow_t ri2(5); ri2 << 3, 1, 0, 4, 2;

	EXPECT_TRUE( test_equal(b1, r1) );
	EXPECT_TRUE( test_equal(i1, ri1) );
	EXPECT_TRUE( test_equal(b2, r2) );
	EXPECT_TRUE( test_equal(i2, ri2) );
}




